# Face detection node for DJI Tello
# 2024/03/14

import cv2
from cv_bridge import CvBridge, CvBridgeError
import math
import rclpy
from rclpy.node import Node
from sensor_msgs.msg import Image
from geometry_msgs.msg import Twist


# 定义人脸检测节点
class FaceDetectionNode(Node):
    # 节点初始化
    def __init__(self):
        super().__init__('face_detection_node')
        # 加载人脸检测器
        self.face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')
        self.bridge = CvBridge()
        # 订阅camera/image_color话题
        self.subscription = self.create_subscription(
            Image,
            'image_raw',
            self.image_callback,
            10)
        self.subscription  # prevent unused variable warning
        # 创建cmd_vel话题的publisher
        self.velocity_publisher = self.create_publisher(
            Twist,
            'cmd_vel',
            10
        )
        self.velocity_command = Twist()# 该变量用于存放控制量

    # 定义回调函数（处理接收到的图像）
    def image_callback(self, msg):
        face_count = 0
        try:
            cv_image = self.bridge.imgmsg_to_cv2(msg, "bgr8")# 使用cv_bridge将ROS图像转化为opencv的数据类型
            img_height, img_width, channels = cv_image.shape# 获取画面宽高
            gray_frame = cv2.cvtColor(cv_image, cv2.COLOR_BGR2GRAY)# 
            # 检测人脸
            faces = self.face_cascade.detectMultiScale(gray_frame, scaleFactor=1.3, minNeighbors=5)
            # 在图像上绘制矩形框
            for (x, y, w, h) in faces:
                cv2.rectangle(cv_image, (x, y), (x+w, y+h), (0, 255, 0), 2)
                face_count += 1
            # 显示结果
            cv2.imshow('Detected Faces', cv_image)
            cv2.waitKey(1)
        except CvBridgeError as e:
            print (e)   
        
        # 判断人脸个数，仅在检测到一张人脸时动作
        if face_count > 1:
            print("Multiple faces detected")
            self.velocity_command.linear.x = 0.0
            self.velocity_command.linear.z = 0.0
            self.velocity_command.angular.z = 0.0
        elif face_count == 1:
            print("Face detected")
            # 计算控制量
            # x轴（前/后）根据人脸相对画面的大小控制；z轴（上/下）根据人脸中心在图像的上下位置控制；
            # z轴角速度（水平旋转）根据人脸中心在图像的左右位置控制
            self.velocity_command.linear.x = 0.3 * ((0.2 * math.sqrt(img_height ** 2 + img_width ** 2) / math.sqrt(h ** 2 + w ** 2)) - 1.0)
            self.velocity_command.linear.z = 1.0 * (0.5 - ((y + h / 2.0) / img_height))
            self.velocity_command.angular.z = 1.0 * (0.5 - ((x + w / 2.0) / img_width))
        else:
            print("No face detected")
            self.velocity_command.linear.x = 0.0
            self.velocity_command.linear.z = 0.0
            self.velocity_command.angular.z = 0.0
        # 发布控制指令消息
        self.velocity_publisher.publish(self.velocity_command)


def main(args=None):
    try:
        rclpy.init(args=args)
        face_detection_node = FaceDetectionNode()
        rclpy.spin(face_detection_node)
    except KeyboardInterrupt:# 关闭窗口
        print("Shutting down face detection node.")
        cv2.destroyAllWindows()


if __name__ == '__main__':
    main()
